import torch.nn
import torchvision
from torch.nn import MaxPool2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

# 数据集
test_dataset = torchvision.datasets.CIFAR10(root='./Dataset',
                                            train=False,
                                            transform=torchvision.transforms.ToTensor())

test_loader = DataLoader(dataset=test_dataset, batch_size=64)

# 最大池化
class Pool(torch.nn.Module):
    def __init__(self):
        super(Pool, self).__init__()
        self.maxpool1 = MaxPool2d(kernel_size=3, ceil_mode=True)

    def forward(self, x):
        x = self.maxpool1(x)
        return x

Pool = Pool()

writer = SummaryWriter("../logs_nn_MaxPool2d")

i = 0
for data in test_loader:
    images, labels = data
    writer.add_images("input", images, i)
    outputs = Pool(images)
    writer.add_images("output", outputs, i)
    i = i + 1

writer.close()